import os

import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader, ConcatDataset
from torchvision import transforms

from yms_class.models.AlexNet import Alex


class CustomDataset(Dataset):
    def __init__(self, root_dir, transform, class_to_label=None):
        self.root_dir = root_dir
        self.transform = transform or transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        self.class_to_label = class_to_label if class_to_label is not None else {}
        self.images = [f for f in os.listdir(root_dir) if f.endswith(('.bmp', '.jpg', '.png'))]

        # 如果没有提供class_to_label字典，我们在这里创建它
        if not self.class_to_label:
            self._create_class_to_label_mapping()
            self.idx_to_labels = {i: cls_name for i, cls_name in enumerate(self.class_to_label)}

    def _create_class_to_label_mapping(self):
        # 假设类别是从0开始编号的连续整数
        self.classes = sorted(set([filename.split('_')[0] for filename in self.images]))
        self.class_to_label = {cls: i for i, cls in enumerate(self.classes)}

    def get_class_to_label(self):
        return self.class_to_label

    def __len__(self):
        return len(self.images)

    def save_label_mapping(self, save_dir=None):
        if save_dir is None:
            save_dir = os.path.join(os.path.dirname(self.root_dir))
        np.save(os.path.join(save_dir, 'idx_to_labels.npy'), self.idx_to_labels)
        np.save(os.path.join(save_dir, 'classes_to_idx.npy'), self.class_to_label)

    def __getitem__(self, idx):
        # 获取图片路径
        image_path = os.path.join(self.root_dir, self.images[idx])
        # 打开图片并转换为RGB格式
        image = Image.open(image_path).convert('RGB')
        # image = Image.open(image_path)
        # 如果有变换，则进行变换
        if self.transform:
            image = self.transform(image)

        # 提取文件名中的类别
        base_filename = os.path.splitext(self.images[idx])[0]
        class_name = base_filename.split('_')[0]
        # 将类别转换为标签
        label = self.class_to_label[class_name]

        return image, label


# class DualImageDataset(Dataset):
#     def __init__(self, root_dir1, root_dir2, transform):
#         self.root_dir1 = root_dir1
#         self.root_dir2 = root_dir2
#         self.transform = transform
#         self.samples = []
#         self.class_to_label = np.load(os.path.join(os.path.dirname(self.root_dir1), 'classes_to_idx.npy'),
#                                       allow_pickle=True).item()
#         self.idx_to_labels = {i: cls_name for i, cls_name in enumerate(self.class_to_label)}
#
#         # 获取两个文件夹中的所有文件（假设都是图片）
#         self.files1 = os.listdir(root_dir1)
#         self.files2 = os.listdir(root_dir2)
#         # 按文件名前缀配对
#         self.samples = self._pair_files()
#
#
#     def _pair_files(self):
#         """根据文件名前缀匹配文件对"""
#         paired_samples = []
#
#         # 创建第二个文件夹的文件名前缀映射
#         file2_prefix_map = {f.split('.')[0]: f for f in self.files2}
#
#         # 遍历第一个文件夹的文件，寻找匹配的前缀
#         for f1 in self.files1:
#             prefix = f1.split('.')[0]
#             # 检查第二个文件夹中是否有相同前缀的文件
#             if prefix in file2_prefix_map:
#                 f2 = file2_prefix_map[prefix]
#                 paired_samples.append((
#                     os.path.join(self.root_dir1, f1),
#                     os.path.join(self.root_dir2, f2),
#                     self.class_to_label[f2.split('_')[0]]
#                 ))
#
#         return paired_samples
#
#     def __len__(self):
#         return len(self.samples)
#
#     def __getitem__(self, idx):
#         """获取一对图片"""
#         path1, path2, label = self.samples[idx]
#
#         # 加载图片（假设都是有效图片格式）
#         img1 = Image.open(path1).convert('RGB')
#         img2 = Image.open(path2).convert('RGB')
#
#         # 应用变换
#         if self.transform:
#             img1 = self.transform(img1)
#             img2 = self.transform(img2)
#
#         return img1, img2, label

class SixImageDataset(Dataset):
    def __init__(self, root_dirs, transform):
        """
        初始化六图片配对数据集
        :param root_dirs: 包含6个文件夹路径的列表，每个文件夹对应一种类型的图片
        :param transform: 应用于图片的变换
        """
        # 确保提供了6个文件夹
        assert len(root_dirs) == 6, "需要提供6个文件夹路径"

        self.root_dirs = root_dirs
        self.transform = transform
        self.samples = []

        # 加载类别映射（假设从第一个文件夹的父目录加载）
        self.class_to_label = np.load(
            os.path.join(os.path.dirname(self.root_dirs[0]), 'classes_to_idx.npy'),
            allow_pickle=True
        ).item()
        self.idx_to_labels = {i: cls_name for i, cls_name in enumerate(self.class_to_label)}
        self.classes = list(self.class_to_label.keys())
        # 获取所有文件夹中的文件并进行配对
        self.samples = self._pair_files()

    def _pair_files(self):
        """根据文件名前缀匹配6张图片"""
        paired_samples = []

        # 从第一个文件夹获取所有文件名前缀作为基准
        first_dir_files = os.listdir(self.root_dirs[0])
        prefixes = [f.split('.')[0] for f in first_dir_files]

        # 为其他5个文件夹创建文件名前缀映射
        dir_prefix_maps = []
        for i in range(1, 6):
            files = os.listdir(self.root_dirs[i])
            prefix_map = {f.split('.')[0]: f for f in files}
            dir_prefix_maps.append(prefix_map)

        # 寻找在所有6个文件夹中都有对应前缀的文件
        for prefix in prefixes:
            # 检查当前前缀在其他文件夹中是否都存在
            has_all = True
            for map in dir_prefix_maps:
                if prefix not in map:
                    has_all = False
                    break

            if has_all:
                # 收集6张图片的路径
                image_paths = [os.path.join(self.root_dirs[0], f"{prefix}{os.path.splitext(first_dir_files[0])[1]}")]
                for i in range(1, 6):
                    filename = dir_prefix_maps[i - 1][prefix]
                    image_paths.append(os.path.join(self.root_dirs[i], filename))

                # 从文件名提取类别标签（假设从第一个文件夹的文件名提取）
                label = self.class_to_label[prefix.split('_')[0]]

                paired_samples.append((image_paths, label))

        return paired_samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        """获取一组6张图片和对应的标签"""
        image_paths, label = self.samples[idx]

        # 加载6张图片
        images = []
        for path in image_paths:
            img = Image.open(path)
            if self.transform:
                img = self.transform(img)
            images.append(img)

        return images[0], images[1], images[2], images[3], images[4], images[5], label



def create_dataloaders(data_path, batch_size, transform=transforms.ToTensor(), num_workers=0, train_shuffle=True):
    # 训练集数据加载器
    train_dir = os.path.join(data_path, 'train')
    train_dataset = CustomDataset(root_dir=train_dir, transform=transform)
    # 初始化验证集Dataset
    validation_dir = os.path.join(data_path, 'val')  # 替换为你的验证集图片目录
    validation_dataset = CustomDataset(root_dir=validation_dir, transform=transform)
    train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=train_shuffle,
                                  num_workers=num_workers)
    val_loader = DataLoader(dataset=validation_dataset, batch_size=batch_size, shuffle=False)
    return train_loader, val_loader

if __name__ == '__main__':
    # data = CustomDataset(root_dir=r'D:\Code\0-data\7-images\2-wear\Nscales(0-2)\val', transform=transforms.Compose([]))
    # data.save_label_mapping()
    img_path = [r'D:\Code\0-data\7-images\2-wear\dataset\scale0_images',
                r'D:\Code\0-data\7-images\2-wear\dataset\scale1_images',
                r'D:\Code\0-data\7-images\2-wear\dataset\scale2_images',
                r'D:\Code\0-data\7-images\2-wear\dataset\scale3_images',
                r'D:\Code\0-data\7-images\2-wear\dataset\scale4_images',
                r'D:\Code\0-data\7-images\2-wear\dataset\scale5_images',
                ]
    data = SixImageDataset(img_path,
                            transform=transforms.ToTensor())
    train = DataLoader(data, batch_size=10)
    model = Alex(num_classes=4)
    for img, labels in train:
        out = model(img)
        print(out)
